package org.helgoboss.scala_additions

import java.io.File
import collection.mutable.Stack

object FileTreeIterator {

    /**
     * Represents a node in the file tree.
     */
    trait Node {
        /**
         * Currently visited file or directory.
         */
        def file: File
        
        /**
         * Path of currently visited file or directory relative to the root directory.
         */
        def path: Path
        
        /**
         * Current directory stack (without the currently processed file/directory). First element ist root directory, last element is the currently walked directory.
         */
        def levels: List[Level]
        
        /**
         * Sets whether iterator shall recurse deeper or not.
         */
        def recurse_=(status: Boolean)
        
        
        /**
         * Returns whether iterator shall recurse deeper or not.
         */
        def recurse: Boolean
        
        /**
         * Sets the walk order for the next recursion.
         */
        def walkOrder_=(walkOrder: Option[(File, File) => Boolean])
        
        /**
         * Returns the walk order for the next recursion.
         */
        def walkOrder: Option[(File, File) => Boolean]
    }
    
    /**
     * Represents a level in the file tree.
     */
    trait Level {
        def dir: File
        def continue_=(status: Boolean)
        def continue: Boolean
    }
}

class FileTreeIterator(root: File, walkOrder: Option[(File, File) => Boolean] = None) extends Iterator[FileTreeIterator.Node] {
    import FileTreeIterator._
    
    private var recurseVar: Boolean = _
    private var walkOrderVar: Option[(File, File) => Boolean] = _
    private var previousNode: Option[Node] = None
    private var nextNode: Option[Node] = None
    private val levelStack = new Stack[InternalLevel]
    
    resetInstructions
    setAsNextNode(root)
                
    def hasNext = {
        determineNextNodeIfNotAlreadyDone
        !nextNode.isEmpty
    }
    
    def next = {
        determineNextNodeIfNotAlreadyDone
        if (nextNode.isEmpty) {
            null
        } else {
            previousNode = nextNode
            nextNode = None
            previousNode.get
        }
    }
    
    
    private def determineNextNodeIfNotAlreadyDone {
        if (nextNode.isEmpty) {
            determineNextNode
        }
    } 
    
    
    private def determineNextNode {
        var stackCompletelyProcessed = false
        
        /* Recurse deeper if recursing is enabled */
        
        if (previousNode.get.file.isDirectory && recurseVar) {
            levelStack push new InternalLevel {
                val dir = previousNode.get.file
                val childrenIterator = {
                    val children = dir.listFiles
                    walkOrderVar match {
                        case Some(x) => children.sortWith(x).iterator
                        case None => children.iterator
                    }
                }
            }
        }
        
        /* Search next file in current level. If no more file is left on this level, go one level up, and so on. */
        
        do {
            if (levelStack.isEmpty) {
                /* Complete iteration finished. */
                stackCompletelyProcessed = true
            } else {
                /* There's something on the level stack to be processed */
                val currentLevel = levelStack.top
                val childrenIterator = currentLevel.childrenIterator
                if (currentLevel.continue && childrenIterator.hasNext) {
                    /* Not all directory entries returned yet. Return next file in current directory. */
                    val file = childrenIterator.next
                    if (file.getName == "..") {
                        /* Ignore */
                    } else {
                        setAsNextNode(file)
                    }
                } else {
                    /* All directory entries processed. Go to parent directory */
                    levelStack.pop
                }
            }
        } while (nextNode.isEmpty && !stackCompletelyProcessed)
    }
    
    private def setAsNextNode(nextFile: File) {
        val node = new Node {
            val levels = levelStack.toList.reverse
            
            val file = nextFile
            
            val path = Path(
                if (levels.isEmpty) {
                    /* This is the root. Occurs only at the first iteration. Its path expression is "" */
                    Nil
                } else if (levels.size == 1) {
                    /* This is a direct sub file of the root. Its path expression is "x" */
                    List(file.getName)
                } else {
                    /* This is a transitive sub file of the root. Its path expression is "x/..." */
                    levels.tail.map(_.dir.getName).toList :+ file.getName
                }
            )
            
            def recurse_=(status: Boolean) {
                recurseVar = status
            }
            
            def recurse = recurseVar
            
            def walkOrder_=(walkOrder: Option[(File, File) => Boolean]) {
                walkOrderVar = walkOrder
            }
            
            def walkOrder = walkOrderVar
        }
        
        nextNode = Some(node)
    }           
    
    private def resetInstructions {
        recurseVar = true
        walkOrderVar = walkOrder
    }
    

    private trait InternalLevel extends Level {
        var continue = true
        val childrenIterator: Iterator[File]
    }
}
